import rdflib
import torch
import numpy as np
import os
from torch_geometric.data import Data, HeteroData
from rdflib import Graph, URIRef, Literal
from rdflib.namespace import RDF, RDFS

class NSKGProcessor:
    """处理nuScenes Knowledge Graph (nSKG)数据的类"""
    
    def __init__(self, nskg_path, ontology_path=None, use_nSTP=False):
        """
        初始化nSKG处理器
        
        Args:
            nskg_path: nSKG数据路径(.ttl文件)
            ontology_path: 本体文件路径(agent和map本体)
            use_nSTP: 是否使用nSTP格式(PyTorch Geometric格式)
        """
        self.nskg_path = nskg_path
        self.ontology_path = ontology_path
        self.use_nSTP = use_nSTP
        
        # 加载知识图谱
        self.graph = self._load_knowledge_graph()
        
        # 定义常用的命名空间
        self.ns = {
            'nsa': URIRef('http://nuscenes.org/agent#'),
            'nsm': URIRef('http://nuscenes.org/map#'),
            'rdf': RDF,
            'rdfs': RDFS
        }
        
        # 缓存已处理的场景图
        self.scene_graphs = {}
        
    def _load_knowledge_graph(self):
        """加载知识图谱"""
        g = Graph()
        
        # 加载nSKG数据
        if os.path.exists(self.nskg_path):
            g.parse(self.nskg_path, format='turtle')
        else:
            raise FileNotFoundError(f"nSKG文件未找到: {self.nskg_path}")
        
        # 加载本体(如果提供)
        if self.ontology_path:
            agent_onto = os.path.join(self.ontology_path, 'nuScenes_agent_onto.ttl')
            map_onto = os.path.join(self.ontology_path, 'nuScenes_map_onto.ttl')
            
            if os.path.exists(agent_onto):
                g.parse(agent_onto, format='turtle')
            
            if os.path.exists(map_onto):
                g.parse(map_onto, format='turtle')
        
        return g
    
    def get_scene_graph(self, scene_token, sample_token):
        """
        获取特定场景和样本的知识图谱
        
        Args:
            scene_token: 场景标识符
            sample_token: 样本标识符
            
        Returns:
            PyTorch Geometric格式的图数据
        """
        # 检查缓存
        cache_key = f"{scene_token}_{sample_token}"
        if cache_key in self.scene_graphs:
            return self.scene_graphs[cache_key]
        
        # 如果使用nSTP格式，直接加载预处理的PyG文件
        if self.use_nSTP:
            pyg_path = os.path.join(
                os.path.dirname(self.nskg_path),
                'nSTP',
                scene_token,
                f"{sample_token}.pt"
            )
            if os.path.exists(pyg_path):
                graph_data = torch.load(pyg_path)
                self.scene_graphs[cache_key] = graph_data
                return graph_data
        
        # 否则，从RDF图中提取子图
        scene_uri = URIRef(f"http://nuscenes.org/scene#{scene_token}")
        sample_uri = URIRef(f"http://nuscenes.org/sample#{sample_token}")
        
        # 查询与该样本相关的所有实体和关系
        entities, relations = self._extract_subgraph(scene_uri, sample_uri)
        
        # 转换为PyTorch Geometric格式
        graph_data = self._convert_to_pyg(entities, relations)
        
        # 缓存结果
        self.scene_graphs[cache_key] = graph_data
        
        return graph_data
    
    def _extract_subgraph(self, scene_uri, sample_uri):
        """从RDF图中提取子图"""
        entities = {}
        relations = []
        
        # 查询样本中的所有实体
        query = f"""
        SELECT ?entity ?type ?x ?y ?z ?length ?width ?height ?yaw
        WHERE {{
            ?entity <http://nuscenes.org/belongsTo> <{sample_uri}> .
            ?entity a ?type .
            OPTIONAL {{ ?entity <http://nuscenes.org/position/x> ?x }}
            OPTIONAL {{ ?entity <http://nuscenes.org/position/y> ?y }}
            OPTIONAL {{ ?entity <http://nuscenes.org/position/z> ?z }}
            OPTIONAL {{ ?entity <http://nuscenes.org/dimension/length> ?length }}
            OPTIONAL {{ ?entity <http://nuscenes.org/dimension/width> ?width }}
            OPTIONAL {{ ?entity <http://nuscenes.org/dimension/height> ?height }}
            OPTIONAL {{ ?entity <http://nuscenes.org/rotation/yaw> ?yaw }}
        }}
        """
        
        for row in self.graph.query(query):
            entity_uri, type_uri = row.entity, row.type
            entity_id = str(entity_uri).split('#')[-1]
            type_id = str(type_uri).split('#')[-1]
            
            # 提取位置和尺寸信息
            pos = [
                float(row.x) if row.x else 0.0,
                float(row.y) if row.y else 0.0,
                float(row.z) if row.z else 0.0
            ]
            
            dim = [
                float(row.length) if row.length else 0.0,
                float(row.width) if row.width else 0.0,
                float(row.height) if row.height else 0.0
            ]
            
            yaw = float(row.yaw) if row.yaw else 0.0
            
            entities[entity_uri] = {
                'id': entity_id,
                'type': type_id,
                'pos': pos,
                'dim': dim,
                'yaw': yaw
            }
        
        # 查询实体间的关系
        query = f"""
        SELECT ?subject ?predicate ?object
        WHERE {{
            ?subject <http://nuscenes.org/belongsTo> <{sample_uri}> .
            ?object <http://nuscenes.org/belongsTo> <{sample_uri}> .
            ?subject ?predicate ?object .
            FILTER(?subject != ?object)
        }}
        """
        
        for row in self.graph.query(query):
            subj, pred, obj = row.subject, row.predicate, row.object
            
            # 只保留我们感兴趣的关系
            pred_str = str(pred)
            if any(ns in pred_str for ns in ['nsa', 'nsm', 'spatial', 'semantic']):
                if subj in entities and obj in entities:
                    rel_type = pred_str.split('#')[-1]
                    relations.append((subj, rel_type, obj))
        
        return entities, relations
    
    def _convert_to_pyg(self, entities, relations):
        """将实体和关系转换为PyTorch Geometric格式"""
        if self.use_nSTP:
            # 创建异构图
            data = HeteroData()
            
            # 处理不同类型的节点
            entity_types = set(entity['type'] for entity in entities.values())
            entity_indices = {}  # 映射实体URI到索引
            
            for entity_type in entity_types:
                # 获取该类型的所有实体
                type_entities = [
                    (uri, entity) for uri, entity in entities.items()
                    if entity['type'] == entity_type
                ]
                
                if not type_entities:
                    continue
                
                # 记录索引映射
                for i, (uri, _) in enumerate(type_entities):
                    entity_indices[uri] = (entity_type, i)
                
                # 提取特征
                x = []
                pos = []
                for _, entity in type_entities:
                    # 基本特征: 位置、尺寸、朝向
                    features = entity['pos'] + entity['dim'] + [np.sin(entity['yaw']), np.cos(entity['yaw'])]
                    x.append(features)
                    pos.append(entity['pos'])
                
                # 添加到图中
                data[entity_type].x = torch.tensor(x, dtype=torch.float)
                data[entity_type].pos = torch.tensor(pos, dtype=torch.float)
            
            # 处理不同类型的边
            relation_types = set(rel_type for _, rel_type, _ in relations)
            
            for rel_type in relation_types:
                # 获取该类型的所有关系
                type_relations = [
                    (subj, obj) for subj, r_type, obj in relations
                    if r_type == rel_type
                ]
                
                if not type_relations:
                    continue
                
                # 按源节点和目标节点类型分组
                grouped_relations = {}
                for subj, obj in type_relations:
                    if subj in entity_indices and obj in entity_indices:
                        subj_type, subj_idx = entity_indices[subj]
                        obj_type, obj_idx = entity_indices[obj]
                        key = (subj_type, rel_type, obj_type)
                        
                        if key not in grouped_relations:
                            grouped_relations[key] = ([], [])
                        
                        src_list, dst_list = grouped_relations[key]
                        src_list.append(subj_idx)
                        dst_list.append(obj_idx)
                
                # 添加到图中
                for (src_type, edge_type, dst_type), (src_indices, dst_indices) in grouped_relations.items():
                    if src_indices and dst_indices:
                        edge_index = torch.tensor([src_indices, dst_indices], dtype=torch.long)
                        data[src_type, edge_type, dst_type].edge_index = edge_index
            
            return data
        else:
            # 创建同构图
            # 映射实体URI到索引
            entity_list = list(entities.keys())
            entity_to_idx = {uri: i for i, uri in enumerate(entity_list)}
            
            # 节点特征
            x = []
            pos = []
            node_types = []
            
            for uri in entity_list:
                entity = entities[uri]
                # 基本特征: 位置、尺寸、朝向
                features = entity['pos'] + entity['dim'] + [np.sin(entity['yaw']), np.cos(entity['yaw'])]
                x.append(features)
                pos.append(entity['pos'])
                
                # 节点类型(作为额外特征)
                type_id = entity['type']
                node_types.append(type_id)
            
            # 边索引和类型
            edge_index = [[], []]
            edge_type = []
            
            for subj, rel_type, obj in relations:
                if subj in entity_to_idx and obj in entity_to_idx:
                    src_idx = entity_to_idx[subj]
                    dst_idx = entity_to_idx[obj]
                    
                    edge_index[0].append(src_idx)
                    edge_index[1].append(dst_idx)
                    edge_type.append(rel_type)
            
            # 创建图数据对象
            data = Data(
                x=torch.tensor(x, dtype=torch.float),
                pos=torch.tensor(pos, dtype=torch.float),
                edge_index=torch.tensor(edge_index, dtype=torch.long),
                edge_attr=torch.tensor([[hash(t) % 100] for t in edge_type], dtype=torch.float),
                node_type=node_types,
                edge_type=edge_type
            )
            
            return data